import re

from search_engine_base import AbstractSearchEngine


# 基于Bag of words算法的搜索引擎实现
class BOWSearchEngine(AbstractSearchEngine):
    # 构造函数
    def __init__(self):
        # 调用父类的构造函数
        super(AbstractSearchEngine, self).__init__()

        # 初始化内部索引表
        # 索引表为一个dict, key为语料库ID value为该语料库文本分词后的单词列表
        self.__index_table = {}

    def process_corpus(self, id, text):
        # 构建内部索引表
        self.__index_table[id] = self.parse_text_to_words(text)

    def search(self, query):
        # 获取query的分词列表
        query_words = self.parse_text_to_words(query)
        result = []
        # 遍历索引表进行匹配
        for id, words in self.__index_table.items():
            if self.match_query(query_words, words):
                result.append(id)

        # 返回搜索结果
        return result

    # 对文本text进行分词处理,返回分词后的单词列表
    @staticmethod
    def parse_text_to_words(text):
        # 使用正则表达式去除标点符号和换行符
        text = re.sub(pattern=r'[^\w ]', repl=' ', string=text)
        # 转为小写
        text = text.lower()
        # 采用空格分隔,生成所有单词的列表
        words = text.split(' ')
        # 过滤掉空白单词
        words = filter(None, words)
        # 返回所有单词的set集合
        return set(words)

    # 判断是否语料库的分词列表包含全部的query单词
    @staticmethod
    def match_query(query_words, corpus_words):
        for q in query_words:
            if q not in corpus_words:
                return False
        return True
